from fastapi import FastAPI, Depends, HTTPException, status, Request, File, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm, APIKeyHeader
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from pydantic import BaseModel
from datetime import datetime, timedelta
from jose import JWTError, jwt
import requests
import json
import uuid
import os
import base64
from dotenv import load_dotenv

# 路径配置
current_file_path = os.path.abspath(__file__)
app_dir = os.path.dirname(current_file_path)
static_dir = os.path.join(app_dir, "static")
templates_dir = os.path.join(app_dir, "templates")

# 加载环境变量
load_dotenv()

# 数据库操作导入（使用上面的 database.py）
from .database import (
    init_db, 
    get_user_by_username, 
    get_user_by_id,
    create_user,
    update_user_avatar,
    get_all_users,
    is_admin,
    # 聊天相关
    create_chat_session,
    get_chat_sessions,
    update_chat_session_title,
    delete_chat_session,
    save_chat_message,
    get_chat_messages,
    clear_chat_messages,
    # API Key相关
    create_api_key,
    get_api_keys_for_user,
    get_api_key_by_key,
    activate_api_key,
    revoke_api_key,
    get_pending_api_keys,
    pwd_context
)

# 初始化数据库
init_db()

# FastAPI应用
app = FastAPI()

# 静态文件和模板
app.mount("/static", StaticFiles(directory=static_dir), name="static")
templates = Jinja2Templates(directory=templates_dir)

# 跨域配置
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 安全配置
SECRET_KEY = os.getenv("SECRET_KEY", "default-dev-key-change-in-production")
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES = int(os.getenv("ACCESS_TOKEN_EXPIRE_MINUTES", "30"))

# Ollama配置
OLLAMA_URL = os.getenv("OLLAMA_URL", "http://192.168.55.3:11434/api/chat")
# MODEL_NAME = "deepseek-r1:latest"
MODEL_NAME = "qwen3:32b"

# 认证方案
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)

# --------------------------
# Pydantic模型（与前端请求/响应匹配）
# --------------------------
class UserCreate(BaseModel):
    username: str
    password: str
    email: str = None

class UserResponse(BaseModel):
    id: int
    username: str
    email: str = None
    is_admin: bool
    avatar: str = None
    created_at: datetime

class Token(BaseModel):
    access_token: str
    token_type: str

class TokenData(BaseModel):
    username: str = None

# 聊天相关模型
class ChatSessionCreate(BaseModel):
    title: str = "新话题"

class ChatSessionResponse(BaseModel):
    id: int
    title: str
    created_at: datetime
    updated_at: datetime

class ChatMessageRequest(BaseModel):
    session_id: int
    message: str

class ChatMessageResponse(BaseModel):
    id: int
    role: str
    content: str
    created_at: datetime

# API Key相关模型
class APIKeyRequest(BaseModel):
    description: str = "默认API密钥"

class APIKeyResponse(BaseModel):
    id: int
    key: str
    description: str
    is_active: bool
    is_revoked: bool
    created_at: datetime
    activated_at: datetime = None

class AvatarUploadRequest(BaseModel):
    avatar_base64: str

# --------------------------
# 工具函数
# --------------------------
def verify_password(plain_password, hashed_password):
    return pwd_context.verify(plain_password, hashed_password)

def get_password_hash(password):
    return pwd_context.hash(password)

def authenticate_user(username: str, password: str):
    user = get_user_by_username(username)
    if not user or not verify_password(password, user["password_hash"]):
        return False
    return user

def create_access_token(data: dict, expires_delta: timedelta = None):
    to_encode = data.copy()
    expire = datetime.utcnow() + (expires_delta if expires_delta else timedelta(minutes=15))
    to_encode.update({"exp": expire})
    return jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)

# --------------------------
# 依赖函数
# --------------------------
async def get_current_user(token: str = Depends(oauth2_scheme)):
    credentials_exception = HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="无法验证用户凭据",
        headers={"WWW-Authenticate": "Bearer"},
    )
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        username = payload.get("sub")
        if not username:
            raise credentials_exception
        user = get_user_by_username(username)
        if not user:
            raise credentials_exception
        return user
    except JWTError:
        raise credentials_exception

async def get_current_active_api_key(
    api_key: str = Depends(api_key_header),
    current_user: dict = Depends(get_current_user)
):
    if not api_key:
        raise HTTPException(status_code=401, detail="未提供API Key", headers={"WWW-Authenticate": "X-API-Key"})
    api_key_data = get_api_key_by_key(api_key)
    if not api_key_data or api_key_data["user_id"] != current_user["id"]:
        raise HTTPException(status_code=401, detail="无效或不属于当前用户的API Key")
    if api_key_data["is_revoked"] or not api_key_data["is_active"]:
        raise HTTPException(status_code=401, detail="API Key已注销或未激活")
    return api_key_data

async def get_current_admin(current_user: dict = Depends(get_current_user)):
    if not is_admin(current_user["id"]):
        raise HTTPException(status_code=403, detail="需要管理员权限")
    return current_user

# --------------------------
# Ollama 交互函数
# --------------------------
def send_to_ollama(prompt: str, history: list):
    try:
        messages = history + [{"role": "user", "content": prompt}]
        response = requests.post(
            OLLAMA_URL,
            headers={"Content-Type": "application/json"},
            data=json.dumps({"model": MODEL_NAME, "messages": messages, "stream": False})
        )
        response.raise_for_status()
        return response.json()["message"]["content"]
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Ollama 通信错误: {str(e)}")

# --------------------------
# 前端页面渲染
# --------------------------
@app.get("/", include_in_schema=False)
def read_root(request: Request):
    return templates.TemplateResponse("index.html", {"request": request})

# --------------------------
# 用户相关接口
# --------------------------
@app.get("/users/me", response_model=UserResponse)
async def read_users_me(current_user: dict = Depends(get_current_user)):
    """获取当前登录用户信息（含头像）"""
    return current_user

@app.post("/register", response_model=UserResponse)
def register(user: UserCreate):
    """注册新用户"""
    if get_user_by_username(user.username):
        raise HTTPException(status_code=400, detail="用户名已存在")
    password_hash = get_password_hash(user.password)
    return create_user(user.username, password_hash, user.email)

@app.post("/token", response_model=Token)
async def login_for_access_token(form_data: OAuth2PasswordRequestForm = Depends()):
    """登录获取令牌"""
    user = authenticate_user(form_data.username, form_data.password)
    if not user:
        raise HTTPException(status_code=401, detail="用户名或密码不正确", headers={"WWW-Authenticate": "Bearer"})
    access_token = create_access_token(
        data={"sub": user["username"]},
        expires_delta=timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
    )
    return {"access_token": access_token, "token_type": "bearer"}

@app.post("/users/avatar")
async def upload_user_avatar(
    req: AvatarUploadRequest,
    current_user: dict = Depends(get_current_user)
):
    """上传用户头像（Base64格式）"""
    update_user_avatar(current_user["id"], req.avatar_base64)
    return {"status": "success", "message": "头像更新成功"}

# --------------------------
# 聊天会话接口（前端历史话题列表）
# --------------------------
@app.post("/chat/sessions", response_model=ChatSessionResponse)
async def create_session(
    session: ChatSessionCreate,
    current_user: dict = Depends(get_current_user)
):
    """创建新聊天会话"""
    return create_chat_session(current_user["id"], session.title)

@app.get("/chat/sessions", response_model=list[ChatSessionResponse])
async def get_sessions(
    limit: int = 10,
    current_user: dict = Depends(get_current_user)
):
    """获取用户的聊天会话列表（历史话题）"""
    return get_chat_sessions(current_user["id"], limit)

@app.put("/chat/sessions/{session_id}/title")
async def update_session_title(
    session_id: int,
    new_title: str,
    current_user: dict = Depends(get_current_user)
):
    """更新会话标题"""
    # 验证会话归属
    sessions = get_chat_sessions(current_user["id"])
    if not any(s["id"] == session_id for s in sessions):
        raise HTTPException(status_code=403, detail="无权修改此会话")
    update_chat_session_title(session_id, new_title)
    return {"status": "success", "message": "会话标题更新成功"}

@app.delete("/chat/sessions/{session_id}")
async def delete_session(
    session_id: int,
    current_user: dict = Depends(get_current_user)
):
    """删除聊天会话"""
    sessions = get_chat_sessions(current_user["id"])
    if not any(s["id"] == session_id for s in sessions):
        raise HTTPException(status_code=403, detail="无权删除此会话")
    delete_chat_session(session_id)
    return {"status": "success", "message": "会话已删除"}

# --------------------------
# 聊天消息接口
# --------------------------
@app.post("/chat/messages", response_model=ChatMessageResponse)
async def send_message(
    req: ChatMessageRequest,
    current_user: dict = Depends(get_current_user),
    api_key: dict = Depends(get_current_active_api_key)
):
    """发送聊天消息（关联会话）"""
    # 验证会话归属
    sessions = get_chat_sessions(current_user["id"])
    if not any(s["id"] == req.session_id for s in sessions):
        raise HTTPException(status_code=403, detail="无权访问此会话")
    
    # 获取历史消息
    history = get_chat_messages(req.session_id)
    history_ollama = [{"role": m["role"], "content": m["content"]} for m in history]
    
    # 调用Ollama获取响应
    assistant_content = send_to_ollama(req.message, history_ollama)
    
    # 保存用户消息
    save_chat_message(req.session_id, current_user["id"], "user", req.message)
    
    # 保存并返回助手消息
    return save_chat_message(req.session_id, current_user["id"], "assistant", assistant_content)

@app.get("/chat/sessions/{session_id}/messages", response_model=list[ChatMessageResponse])
async def get_messages(
    session_id: int,
    current_user: dict = Depends(get_current_user)
):
    """获取会话的所有消息"""
    sessions = get_chat_sessions(current_user["id"])
    if not any(s["id"] == session_id for s in sessions):
        raise HTTPException(status_code=403, detail="无权访问此会话")
    return get_chat_messages(session_id)

@app.delete("/chat/sessions/{session_id}/messages")
async def clear_messages(
    session_id: int,
    current_user: dict = Depends(get_current_user)
):
    """清空会话的消息"""
    sessions = get_chat_sessions(current_user["id"])
    if not any(s["id"] == session_id for s in sessions):
        raise HTTPException(status_code=403, detail="无权操作此会话")
    clear_chat_messages(session_id)
    return {"status": "success", "message": "聊天消息已清空"}

# --------------------------
# API Key相关接口
# --------------------------
@app.post("/api-keys", response_model=APIKeyResponse)
def create_new_api_key(
    req: APIKeyRequest,
    current_user: dict = Depends(get_current_user)
):
    """用户申请新API Key"""
    api_key_str = f"sk-{str(uuid.uuid4()).replace('-', '')[:16]}"  # 生成16位Key
    return create_api_key(current_user["id"], api_key_str, req.description)

@app.get("/api-keys", response_model=list[APIKeyResponse])
def get_user_api_keys(current_user: dict = Depends(get_current_user)):
    """获取当前用户的所有API Key"""
    return get_api_keys_for_user(current_user["id"])

@app.post("/api-keys/{key_id}/revoke")
def revoke_user_api_key(
    key_id: int,
    current_user: dict = Depends(get_current_user)
):
    """用户注销自己的API Key"""
    # 验证Key归属
    keys = get_api_keys_for_user(current_user["id"])
    if not any(k["id"] == key_id for k in keys):
        raise HTTPException(status_code=403, detail="无权注销此API Key")
    success = revoke_api_key(key_id, user_id=current_user["id"])
    if not success:
        raise HTTPException(status_code=404, detail="API Key不存在")
    return {"status": "success", "message": "API Key已注销"}

# --------------------------
# 管理员接口
# --------------------------
@app.get("/admin/users", response_model=list[UserResponse])
def admin_get_all_users(current_admin: dict = Depends(get_current_admin)):
    """管理员获取所有用户"""
    return get_all_users()

@app.get("/admin/pending-api-keys", response_model=list[dict])
def admin_get_pending_api_keys(current_admin: dict = Depends(get_current_admin)):
    """管理员获取待激活API Key（含用户名）"""
    return get_pending_api_keys()

@app.post("/admin/api-keys/{key_id}/activate")
def admin_activate_api_key(
    key_id: int,
    current_admin: dict = Depends(get_current_admin)
):
    """管理员激活API Key"""
    activated_key = activate_api_key(key_id, current_admin["id"])
    if not activated_key:
        raise HTTPException(status_code=404, detail="API Key不存在或已注销")
    return {"status": "success", "message": "API Key已激活", "key": activated_key}

@app.post("/admin/api-keys/{key_id}/revoke")
def admin_revoke_api_key(
    key_id: int,
    current_admin: dict = Depends(get_current_admin)
):
    """管理员注销API Key"""
    success = revoke_api_key(key_id, admin_id=current_admin["id"])
    if not success:
        raise HTTPException(status_code=404, detail="API Key不存在")
    return {"status": "success", "message": "API Key已注销"}

# --------------------------
# 健康检查
# --------------------------
@app.get("/health")
def health_check():
    try:
        response = requests.get(OLLAMA_URL.replace("/api/chat", ""))
        return {
            "status": "healthy" if response.status_code == 200 else "unhealthy",
            "ollama_available": response.status_code == 200,
            "ollama_url": OLLAMA_URL
        }
    except Exception as e:
        return {"status": "unhealthy", "ollama_available": False, "ollama_url": OLLAMA_URL, "error": str(e)}